#ifndef MODEL_BASIC_DATA_STRUCTURE_H
#define MODEL_BASIC_DATA_STRUCTURE_H

#include <iostream>
#include <vector>
#include "Eigen/Dense"
#include "data_structure.h"
#include "utilities/math.hpp"

using namespace Eigen;
typedef Eigen::Matrix<double, 6, 6> Matrix6d;

namespace ZORAL
{
    class Position : public Vector3d
    {
    private:
    public:
        Position()
        {
            this->setZero();
        }
        explicit Position(const Vector3d &p);
        explicit Position(const std::vector<double> &p);
        explicit Position(const double p[3]);
        Position(const double &x, const double &y, const double &z);

        Position &operator=(const Position &p) = default;
        friend Position operator+(const Position &p1, const Position &p2);
        friend Position operator*(const Position &p, const double &k);
        friend Position operator*(const double &k, const Position &p);
    };

    class Rotation
    {
    private:
        double R_[9] = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0};

    public:
        Rotation();
        Rotation(const double &Xx, const double &Xy, const double &Xz, const double &Yx, const double &Yy,
                 const double &Yz, const double &Zx, const double &Zy, const double &Zz);
        explicit Rotation(const std::vector<std::vector<double>> &R);
        explicit Rotation(const Matrix3d &R);
        explicit Rotation(const double R[9]);

        inline double *data()
        {
            return this->R_;
        }
        inline const double *data() const
        {
            return this->R_;
        }
        Rotation transpose() const;
        Matrix3d toMatrix3d() const;
        MatrixXd toMatrixXd() const;

        double trace() const;
        Rotation log3() const;
        Rotation exp3() const;

        double operator[](const int i) const
        {
            return R_[i];
        }
        double &operator[](const int i)
        {
            return R_[i];
        }
        double operator()(const int i) const
        {
            return R_[i];
        }
        double &operator()(const int i)
        {
            return R_[i];
        }
        Rotation &operator=(const Rotation &R) = default;
        friend Rotation operator+(const Rotation &lhs, const Rotation &rhs);
        friend Rotation operator*(const Rotation &lhs, const Rotation &rhs);
        friend MatrixXd operator*(const Rotation &lhs, const MatrixXd &rhs);
        friend MatrixXd operator*(const MatrixXd &lhs, const Rotation &rhs);
        friend Position operator*(const Rotation &R, const Position &p);
        friend Rotation operator*(const Rotation &R, const double &k);
        friend Rotation operator*(const double &k, const Rotation &R);
        friend Vector3d operator*(const Rotation &R, const Vector3d &v);
    };

    class Wrench : public Vector6d
    {
    public:
        Wrench()
        {
            this->setZero();
        }
        Wrench(const double &Mx, const double &My, const double &Mz, const double &Fx, const double &Fy,
               const double &Fz);
        Wrench(const Vector3d &torque, const Vector3d &force);
        explicit Wrench(const Vector6d &vec);
        explicit Wrench(const std::vector<double> &vec);
        explicit Wrench(const double vec[6]);

        inline Vector3d m() const
        {
            return this->head(3);
        }
        inline Vector3d f() const
        {
            return this->tail(3);
        }
    };

    class Twist : public Vector6d
    {
    public:
        Twist()
        {
            this->setZero();
        }
        Twist(const Vector3d &w, const Vector3d &v)
        {
            this->head(3) = w;
            this->tail(3) = v;
        }
        explicit Twist(const std::vector<double> &vec)
        {
            memcpy(this->data(), vec.data(), sizeof(double) * WORKSPACE_DIM);
        }
        explicit Twist(const Vector6d &vec)
        {
            memcpy(this->data(), vec.data(), sizeof(double) * WORKSPACE_DIM);
        }
        Twist(const double &wx, const double &wy, const double &wz, const double &vx, const double &vy,
              const double &vz);

        inline Vector3d w() const
        {
            return this->head(3);
        }
        inline Vector3d v() const
        {
            return this->tail(3);
        }
        Twist transpose(const Twist &v) const;

        Twist cross(const Twist &t) const;
        Wrench cross(const Wrench &w) const;

        friend Twist operator+(const Twist &lhs, const Twist &rhs);
        friend Twist operator*(const double &lhs, const Twist &rhs);
    };

    class HomoTrans
    {
    private:
        Position p_;
        Rotation R_;

    public:
        HomoTrans();
        HomoTrans(const Rotation &R, const Position &p);
        explicit HomoTrans(const std::vector<double> &T);
        explicit HomoTrans(const Matrix4d &T);
        inline Rotation R() const
        {
            return R_;
        }
        inline Rotation R(const Rotation &rhs)
        {
            R_ = rhs;
            return R_;
        }
        inline Position p() const
        {
            return p_;
        }
        inline Position p(const Position &rhs)
        {
            p_ = rhs;
            return p_;
        }

        friend HomoTrans operator*(const HomoTrans &lhs, const HomoTrans &rhs);
        friend HomoTrans operator*(const HomoTrans &T, const double &k);
        friend HomoTrans operator*(const double &k, const HomoTrans &T);
        friend Twist operator*(const HomoTrans &trans, const Twist &twist);
        friend Wrench operator*(const HomoTrans &trans, const Wrench &wrench);

        HomoTrans inv();
        Wrench transpose(const Wrench &F) const;
        inline void setIdentity()
        {
            R_ = Rotation();
            p_ = Position();
        }

        inline void set(const Rotation &R, const Position &p)
        {
            R_ = R;
            p_ = p;
        }

        inline void setR(const Rotation &R)
        {
            R_ = R;
        }

        inline void setP(const Position &p)
        {
            p_ = p;
        }

        Matrix4d toMatrix4d() const;
    };

    class SpaceInertia
    {
    private:
        Matrix6d I_;
        Matrix3d Ic_;
        Vector3d c_;
        double m_;
        void calSpaceInertia();

    public:
        SpaceInertia()
        {
            I_ = Matrix6d();
            Ic_ = Matrix3d();
            c_ = Vector3d();
            m_ = 0.0;
        }
        SpaceInertia(const double &m, const double &cx, const double &cy, const double &cz, const double &Ixx,
                     const double &Iyy, const double &Izz, const double &Ixy, const double &Ixz, const double &Iyz);
        SpaceInertia(const double &m, const Vector3d &c, const Matrix3d &Ic);
        SpaceInertia(const double &m, const Vector3d &c, const Vector6d &Ic);
        explicit SpaceInertia(const std::vector<double> &data);

        friend Wrench operator*(const SpaceInertia &lhs, const Twist &rhs);
    };

    class Jacobian : public MatrixXd
    {
    public:
        Jacobian() = delete;
        explicit Jacobian(size_t columns);
        Jacobian(const Jacobian &arg) : MatrixXd(arg)
        {
        }
        explicit Jacobian(const MatrixXd &arg) : MatrixXd(arg)
        {
        }
        Jacobian(const Eigen::MatrixXd &Jw, const Eigen::MatrixXd &Jv);

        void resize(int n);
        Twist getColumn(size_t i);
        void setColumn(size_t i, const Twist &S);
        MatrixXd Jw() const
        {
            return this->topRows(ROT_DIM);
        }
        MatrixXd Jv() const
        {
            return this->bottomRows(TRANS_DIM);
        }

        friend MatrixXd operator*(const Jacobian &jac, const HomoTrans &T);
        friend Jacobian operator*(const HomoTrans &T, const Jacobian &jac);
    };
} // namespace ZORAL
#endif // MODEL_BASIC_DATA_STRUCTURE_H
